#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/time.h>

#include "rtp.h"
#include "util.h"

#define TIME(x) gettimeofday(&x, NULL)
#define DIFF(x, y) ((x.tv_sec - y.tv_sec) * 1000000 + (x.tv_usec - y.tv_usec))

void rcb_init(uint32_t window_size, struct sockaddr_in *receiver_addr)
{
    if (rcb == NULL)
    {
        rcb = (rcb_t *)calloc(1, sizeof(rcb_t));
    }
    else
    {
        perror("The current version of the rtp protocol only supports a single connection");
        exit(EXIT_FAILURE);
    }
    rcb->window_size = window_size;
    rcb->addr_len = sizeof(struct sockaddr_in);
    if (receiver_addr != NULL)
        rcb->receiver_addr = *receiver_addr;
    rcb->seq = 0;
    rcb->waitseq = 0;

    memset(rcb->wlen, -1, sizeof(rcb->wlen));
    // TODO: you can initialize your RTP-related fields here
}

void outp_rtp(rtp_header_t *rtp)
{
    printf("    Type: %u\n    Seq_num: %u\n    Length: %u\n    Checksum: %u\n", rtp->type, rtp->seq_num, rtp->length, rtp->checksum);
}

char check(rtp_header_t *rtp, int recv_bytes, u_int8_t chk_type)
{
    if (recv_bytes < sizeof(rtp_header_t))
        return 0;
    if (rtp->type != chk_type)
        return 0;
    if (rtp->type != RTP_DATA && rtp->length > 0)
        return 0;
    if (recv_bytes != rtp->length + sizeof(rtp_header_t))
        return 0;
    uint32_t pkt_checksum = rtp->checksum;
    rtp->checksum = 0;
    uint32_t computed_checksum = compute_checksum(rtp, recv_bytes);
    rtp->checksum = pkt_checksum;
    if (pkt_checksum == computed_checksum)
        return 1;
    return 0;
}

void cons(rtp_header_t *rtp, u_int8_t type, u_int16_t length, u_int32_t seq_num)
{
    rtp->type = type;
    rtp->checksum = 0;
    rtp->length = length;
    rtp->seq_num = seq_num;
    rtp->checksum = compute_checksum((void *)rtp, length + sizeof(rtp_header_t));
}

int Sendto(int sockfd, const void *msg, int len, int flags, const struct sockaddr *to, socklen_t tolen)
{
    int sent_bytes = sendto(sockfd, msg, len, flags, to, tolen);
    if (sent_bytes != len)
    {
        perror("send error");
        exit(EXIT_FAILURE);
    }
    return sent_bytes;
}

int Recvfrom(int sockfd, void *buf, int len, int flags, struct sockaddr *from, socklen_t *fromlen)
{
    int recv_bytes = recvfrom(sockfd, buf, len, flags, from, fromlen);
    if (recv_bytes < 0)
    {
        perror("receive error"),
            exit(EXIT_FAILURE);
    }
    return recv_bytes;
}

/*********************** Note ************************/
/* RTP in Assignment 2 only supports single connection.
/* Therefore, we can initialize the related fields of RTP when creating the socket.
/* rcb is a global varialble, you can directly use it in your implementatyion.
/*****************************************************/
int rtp_socket(uint32_t window_size, struct sockaddr_in *receiver_addr)
{
    rcb_init(window_size, receiver_addr);
    // create UDP socket
    return socket(AF_INET, SOCK_DGRAM, 0);
}

int rtp_bind(int sockfd, struct sockaddr *addr, socklen_t addrlen)
{
    return bind(sockfd, addr, addrlen);
}

int rtp_listen(int sockfd, int backlog)
{
    // TODO: listen for the START message from sender and send back ACK
    // In standard POSIX API, backlog is the number of connections allowed on the incoming queue.
    // For RTP, backlog is always 1
    rtp_header_t r_rtp;
    for (char ok = 0; !ok;)
    {
        int recv_bytes = Recvfrom(sockfd, (void *)&r_rtp, sizeof(rtp_header_t), 0, (struct sockaddr *)&rcb->sender_addr, &rcb->addr_len);
        ok = check(&r_rtp, recv_bytes, RTP_START);
        if (ok == 0)
        {
            puts("[INFO] : START check fail.");
            return 0;
        }
    }
    puts("[INFO] : START is received");
    rtp_header_t s_rtp;
    cons(&s_rtp, RTP_ACK, 0, r_rtp.seq_num);
    // printf("%s:%d %d\n", inet_ntoa(rcb->sender_addr.sin_addr), ntohs(rcb->sender_addr.sin_port), rcb->addr_len);
    Sendto(sockfd, (void *)&s_rtp, sizeof(rtp_header_t), 0, (struct sockaddr *)&rcb->sender_addr, rcb->addr_len);
    printf("[INFO] : Listened successfully at %s:%d\n", inet_ntoa(rcb->sender_addr.sin_addr), ntohs(rcb->sender_addr.sin_port));
    return 0;
}

int rtp_accept(int sockfd, struct sockaddr *addr, socklen_t *addrlen)
{
    // Since RTP in Assignment 2 only supports one connection,
    // there is no need to implement accpet function.
    // You don’t need to make any changes to this function.
    addr = (struct sockaddr *)&rcb->sender_addr;
    addrlen = &rcb->addr_len;
    return 1;
}

int rtp_connect(int sockfd, const struct sockaddr *addr, socklen_t addrlen)
{
    // TODO: send START message and wait for its ACK
    // rcb->receiver_addr = *addr;

    rtp_header_t s_rtp,
        r_rtp;
    char ack = 0;
    int recv_bytes;
    struct timeval start_t, cur_t;

    cons(&s_rtp, RTP_START, 0, 666623333);

    Sendto(sockfd, (void *)&s_rtp, sizeof(rtp_header_t), 0, addr, addrlen);
    for (TIME(start_t), TIME(cur_t); !ack && DIFF(cur_t, start_t) <= 500000; TIME(cur_t))
    {
        recv_bytes = recvfrom(sockfd, (void *)&r_rtp, sizeof(rtp_header_t), MSG_DONTWAIT, NULL, 0);
        ack = (check(&r_rtp, recv_bytes, RTP_ACK) && r_rtp.seq_num == s_rtp.seq_num);
    }
    if (ack)
    {
        return 0;
    }
    puts("[INFO] : Waiting for ACK for START timeout.");

    cons(&s_rtp, RTP_END, 0, 998244353);
    ack = 0;
    Sendto(sockfd, (void *)&s_rtp, sizeof(rtp_header_t), 0, addr, addrlen);
    puts("[INFO] : Send END");
    // puts("[INFO] : Waiting for ACK for END");
    for (TIME(start_t), TIME(cur_t); !ack && DIFF(cur_t, start_t) <= 500000; TIME(cur_t))
    {
        recv_bytes = recvfrom(sockfd, (void *)&r_rtp, sizeof(rtp_header_t), MSG_DONTWAIT, NULL, 0);
        ack = (check(&r_rtp, recv_bytes, RTP_ACK) && r_rtp.seq_num == s_rtp.seq_num);
    }
    if (!ack)
        puts("[INFO] : Waiting for ACK for END timeout.");
    return -1;
}

int rtp_close_receiver(int sockfd)
{
    free(rcb);
    return close(sockfd);
}

int rtp_close_sender(int sockfd)
{
    rtp_header_t s_rtp, r_rtp;
    char ack;
    struct timeval start_t, cur_t;
    int recv_bytes;

    cons(&s_rtp, RTP_END, 0, 998244353);
    Sendto(sockfd, (void *)&s_rtp, sizeof(rtp_header_t), 0, (struct sockaddr *)&rcb->receiver_addr, rcb->addr_len);
    puts("[INFO] : Send END");

    ack = 0;
    for (TIME(start_t), TIME(cur_t); !ack && DIFF(cur_t, start_t) <= 500000; TIME(cur_t))
    {
        recv_bytes = recvfrom(sockfd, (void *)&r_rtp, sizeof(rtp_header_t), MSG_DONTWAIT, NULL, 0);
        ack = (check(&r_rtp, recv_bytes, RTP_ACK) && r_rtp.seq_num == s_rtp.seq_num);
    }
    if (!ack)
        puts("[INFO] : Waiting for ACK for END timeout.");

    free(rcb);
    puts("Closed.");
    return 0;
}

int rtp_sendto(int sockfd, const void *msg, int len, int flags, const struct sockaddr *to, socklen_t tolen)
{
    // TODO: send message

    char buffer[BUFFER_SIZE];
    rtp_header_t r_rtp, s_rtp;
    struct timeval start_t, cur_t;
    int recv_bytes;

    TIME(start_t);
    for (int offset = 0; offset < len || rcb->seq > rcb->waitseq;)
    {
        if (offset < len && rcb->seq - rcb->waitseq < rcb->window_size)
        {
            int plen = PKT_SIZE < len - offset ? PKT_SIZE : len - offset;

            rtp_header_t *rtp = (rtp_header_t *)buffer;
            memcpy((void *)buffer + sizeof(rtp_header_t), msg + offset, plen);
            cons(rtp, RTP_DATA, plen, rcb->seq);

            Sendto(sockfd, buffer, sizeof(rtp_header_t) + plen, 0, to, tolen);
            printf("[INFO] : Send (seq %d, waitseq %d)\n", rcb->seq, rcb->waitseq);

            rcb->seq++;
            offset += PKT_SIZE;
        }
        for (TIME(cur_t); DIFF(cur_t, start_t) <= 500000; TIME(cur_t))
        {
            recv_bytes = recvfrom(sockfd, (void *)&r_rtp, sizeof(rtp_header_t), MSG_DONTWAIT, NULL, 0);
            if (recv_bytes <= 0)
                break;
            if (check(&r_rtp, recv_bytes, RTP_ACK))
            {
                printf("[INFO] : Recieve ACK %d (seq %d, waitseq %d)\n", r_rtp.seq_num, rcb->seq, rcb->waitseq);
                if (r_rtp.seq_num > rcb->waitseq)
                {
                    rcb->waitseq = r_rtp.seq_num;
                    TIME(start_t);
                    printf("[INFO] : Move Window.\n");
                }
            }
        }
        if (DIFF(cur_t, start_t) > 500000)
        {
            TIME(start_t);
            offset -= (rcb->seq - rcb->waitseq) * PKT_SIZE;
            rcb->seq = rcb->waitseq;
            printf("[INFO] : Retransmit. (seq %d, waitseq %d)\n", rcb->seq, rcb->waitseq);
        }
    }

    return 1;
}

int rtp_recvfrom(int sockfd, void *buf, int len, int flags, struct sockaddr *from, socklen_t *fromlen)
{
    // TODO: recv message

    char buffer[BUFFER_SIZE];
    // char window[MAX_WINDOWS][PKT_SIZE]; // payload of each window
    char(*window)[PKT_SIZE];
    window = rcb->window;
    int *wlen; // payload len for each window
    wlen = rcb->wlen;
    memset(wlen, -1, sizeof(wlen));
    int curlen, recv_bytes;
    struct timeval start_t, cur_t;

    for (curlen = 0; curlen + PKT_SIZE <= len;)
    {
        for (TIME(start_t), TIME(cur_t), recv_bytes = 0; recv_bytes <= 0 && DIFF(cur_t, start_t) <= 500000; TIME(cur_t))
            recv_bytes = recvfrom(sockfd, buffer, BUFFER_SIZE, MSG_DONTWAIT, from, fromlen);
        if (recv_bytes <= 0)
            break;

        rtp_header_t *rtp = (rtp_header_t *)buffer;
        rtp_header_t s_rtp;

        if (check(rtp, recv_bytes, RTP_END))
        {
            puts("[INFO] : Recieve END");
            cons(&s_rtp, RTP_ACK, 0, rtp->seq_num);
            Sendto(sockfd, (void *)&s_rtp, sizeof(rtp_header_t), 0, (struct sockaddr *)&rcb->sender_addr, rcb->addr_len);
            puts("[INFO] : Send ACK for END");
            puts("Closed.");
            break;
        }
        if (check(rtp, recv_bytes, RTP_DATA))
        {
            printf("[INFO] : Recieve Packet %d (waitseq %d)\n", rtp->seq_num, rcb->waitseq);

            if (rtp->seq_num <= rcb->waitseq + rcb->window_size && rtp->seq_num >= rcb->waitseq)
            {
                int w = rtp->seq_num % rcb->window_size;
                wlen[w] = rtp->length;
                outp_rtp(rtp);
                memcpy(window[w], buffer + sizeof(rtp_header_t), wlen[w]);
                printf("[INFO] : Packet %d stored in buffer\n", rtp->seq_num);
            }
            else
            {
                printf("[INFO] : Packet %d out of window\n", rtp->seq_num);
            }

            if (rtp->seq_num == rcb->waitseq)
            {
                int i;
                for (i = rcb->waitseq; i < rcb->waitseq + rcb->window_size; ++i)
                {
                    int w = i % rcb->window_size;
                    if (wlen[w] == -1)
                        break;
                    memcpy(buf + curlen, window[w], wlen[w]);
                    curlen += wlen[w];
                    wlen[w] = -1;
                }
                rcb->waitseq = i;
            }
            cons(&s_rtp, RTP_ACK, 0, rcb->waitseq);

            Sendto(sockfd, (void *)&s_rtp, sizeof(rtp_header_t), 0, (struct sockaddr *)&rcb->sender_addr, rcb->addr_len);
            printf("[INFO] : Send ACK (expected for %d)\n", rcb->waitseq);
        }
    }

    return curlen;
}

int rtp_sendto_opt(int sockfd, const void *msg, int len, int flags, const struct sockaddr *to, socklen_t tolen)
{
    // TODO: send message

    char buffer[BUFFER_SIZE], ack[MAX_WINDOWS];
    rtp_header_t r_rtp, s_rtp;
    struct timeval start_t, cur_t;
    int recv_bytes;

    memset(ack, 0, sizeof(ack));
    TIME(start_t);
    for (int offset = 0; offset < len || rcb->seq > rcb->waitseq;)
    {
        if (ack[rcb->seq % rcb->window_size])
        {
            rcb->seq++;
            offset += PKT_SIZE;
            continue;
        }
        if (offset < len && rcb->seq - rcb->waitseq < rcb->window_size)
        {
            int plen = PKT_SIZE < len - offset ? PKT_SIZE : len - offset;

            rtp_header_t *rtp = (rtp_header_t *)buffer;
            memcpy((void *)buffer + sizeof(rtp_header_t), msg + offset, plen);
            cons(rtp, RTP_DATA, plen, rcb->seq);

            Sendto(sockfd, buffer, sizeof(rtp_header_t) + plen, 0, to, tolen);
            printf("[INFO] : Send (seq %d, waitseq %d)\n", rcb->seq, rcb->waitseq);

            rcb->seq++;
            offset += PKT_SIZE;
        }
        for (TIME(cur_t); DIFF(cur_t, start_t) <= 500000; TIME(cur_t))
        {
            recv_bytes = recvfrom(sockfd, (void *)&r_rtp, sizeof(rtp_header_t), MSG_DONTWAIT, NULL, 0);
            if (recv_bytes <= 0)
                break;
            if (check(&r_rtp, recv_bytes, RTP_ACK))
            {
                int rseq = r_rtp.seq_num;
                printf("[INFO] : Recieve ACK %d (seq %d, waitseq %d)\n", r_rtp.seq_num, rcb->seq, rcb->waitseq);
                // if (r_rtp.seq_num > rcb->waitseq)
                // {
                //     rcb->waitseq = r_rtp.seq_num;
                //     TIME(start_t);
                //     printf("[INFO] : Move Window.\n");
                // }
                if (rseq < rcb->waitseq || rseq >= rcb->waitseq + rcb->window_size)
                    continue;
                ack[rseq % rcb->window_size] = 1;
                for (; rcb->waitseq < rcb->seq && ack[rcb->waitseq % rcb->window_size]; ++rcb->waitseq)
                {
                    ack[rcb->waitseq % rcb->window_size] = 0;
                    TIME(start_t);
                }
                printf("[INFO] : Move Window (seq %d, waitseq %d)\n", rcb->seq, rcb->waitseq);
            }
        }
        if (DIFF(cur_t, start_t) > 500000)
        {
            TIME(start_t);
            offset -= (rcb->seq - rcb->waitseq) * PKT_SIZE;
            rcb->seq = rcb->waitseq;
            printf("[INFO] : Retransmission starts. (waitseq %d)\n", rcb->waitseq);
        }
    }

    return 1;
}

int rtp_recvfrom_opt(int sockfd, void *buf, int len, int flags, struct sockaddr *from, socklen_t *fromlen)
{
    // TODO: recv message

    char buffer[BUFFER_SIZE];
    // char window[MAX_WINDOWS][PKT_SIZE]; // payload of each window
    // int wlen[MAX_WINDOWS];              // payload len for each window
    // memset(wlen, -1, sizeof(wlen));
    char(*window)[PKT_SIZE];
    window = rcb->window;
    int *wlen; // payload len for each window
    wlen = rcb->wlen;
    int curlen, recv_bytes;
    struct timeval start_t, cur_t;

    for (curlen = 0; curlen + PKT_SIZE <= len;)
    {
        for (TIME(start_t), TIME(cur_t), recv_bytes = 0; recv_bytes <= 0 && DIFF(cur_t, start_t) <= 5000000; TIME(cur_t))
            recv_bytes = recvfrom(sockfd, buffer, BUFFER_SIZE, MSG_DONTWAIT, from, fromlen);
        if (recv_bytes <= 0)
            break;

        rtp_header_t *rtp = (rtp_header_t *)buffer;
        rtp_header_t s_rtp;

        if (check(rtp, recv_bytes, RTP_END))
        {
            puts("[INFO] : Recieve END");
            cons(&s_rtp, RTP_ACK, 0, rtp->seq_num);
            Sendto(sockfd, (void *)&s_rtp, sizeof(rtp_header_t), 0, (struct sockaddr *)&rcb->sender_addr, rcb->addr_len);
            puts("[INFO] : Send ACK for END");
            puts("Closed.");
            break;
        }
        if (check(rtp, recv_bytes, RTP_DATA))
        {
            printf("[INFO] : Recieve Packet %d (waitseq %d)\n", rtp->seq_num, rcb->waitseq);

            if (rtp->seq_num <= rcb->waitseq + rcb->window_size && rtp->seq_num >= rcb->waitseq)
            {
                int w = rtp->seq_num % rcb->window_size;
                wlen[w] = rtp->length;
                outp_rtp(rtp);
                memcpy(window[w], buffer + sizeof(rtp_header_t), wlen[w]);
                printf("[INFO] : Packet %d stored in buffer\n", rtp->seq_num);

                cons(&s_rtp, RTP_ACK, 0, rtp->seq_num);
                Sendto(sockfd, (void *)&s_rtp, sizeof(rtp_header_t), 0, (struct sockaddr *)&rcb->sender_addr, rcb->addr_len);
                printf("[INFO] : Send ACK (for %d)\n", rtp->seq_num);
            }
            else
            {
                printf("[INFO] : Packet %d out of window\n", rtp->seq_num);
            }

            if (rtp->seq_num == rcb->waitseq)
            {
                int i;
                for (i = rcb->waitseq; i < rcb->waitseq + rcb->window_size; ++i)
                {
                    int w = i % rcb->window_size;
                    if (wlen[w] == -1)
                        break;
                    memcpy(buf + curlen, window[w], wlen[w]);
                    curlen += wlen[w];
                    wlen[w] = -1;
                }

                rcb->waitseq = i;
            }
        }
    }

    return curlen;
}
